package org.apache.flink.connector.scylladb.sink;

import com.datastax.driver.core.*;
import com.datastax.driver.core.policies.DCAwareRoundRobinPolicy;
import com.datastax.driver.core.policies.DefaultRetryPolicy;
import com.datastax.driver.core.policies.TokenAwarePolicy;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.connector.scylladb.serialization.ScyllaDBRowDataSerializer;
import org.apache.flink.streaming.api.functions.sink.RichSinkFunction;
import org.apache.flink.table.data.RowData;
import org.apache.flink.table.types.logical.RowType;
import org.apache.flink.types.RowKind;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.net.ssl.SSLContext;
import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * ScyllaDB 目标函数，负责将数据写入 ScyllaDB
 */
public class ScyllaDBSinkFunction extends RichSinkFunction<RowData> {

    private static final Logger LOG = LoggerFactory.getLogger(ScyllaDBSinkFunction.class);

    private final String url;
    private final String keySpace;
    private final String tableName;
    private final String username;
    private final String password;
    private final int bufferFlushMaxRows;
    private final Duration bufferFlushInterval;
    private final int maxRetries;
    private final int connectionPoolLocalSize;
    private final int connectionPoolRemoteSize;
    private final Duration connectionTimeout;
    private final Duration requestTimeout;
    private final String consistencyLevel;
    private final boolean sslEnabled;
    private final String sslKeystorePath;
    private final String sslKeystorePassword;
    private final String sslTruststorePath;
    private final String sslTruststorePassword;
    private final RowType rowType;
    private final Set<String> appendSetFields;
    private final List<String> primaryKeys;

    private transient Cluster cluster;
    private transient Session session;
    private transient ScyllaDBRowDataSerializer serializer;
    private transient PreparedStatement insertStatement;
    private transient PreparedStatement updateStatement;
    private transient PreparedStatement deleteStatement;
    private transient PreparedStatement appendSetStatement;
    private transient List<RowData> buffer;
    private transient ScheduledExecutorService scheduler;

    public ScyllaDBSinkFunction(
            String url,
            String keySpace,
            String tableName,
            String username,
            String password,
            int bufferFlushMaxRows,
            Duration bufferFlushInterval,
            int maxRetries,
            int connectionPoolLocalSize,
            int connectionPoolRemoteSize,
            Duration connectionTimeout,
            Duration requestTimeout,
            String consistencyLevel,
            boolean sslEnabled,
            String sslKeystorePath,
            String sslKeystorePassword,
            String sslTruststorePath,
            String sslTruststorePassword,
            RowType rowType,
            Set<String> appendSetFields,
            List<String> primaryKeys) {
        
        this.url = url;
        this.keySpace = keySpace;
        this.tableName = tableName;
        this.username = username;
        this.password = password;
        this.bufferFlushMaxRows = bufferFlushMaxRows;
        this.bufferFlushInterval = bufferFlushInterval;
        this.maxRetries = maxRetries;
        this.connectionPoolLocalSize = connectionPoolLocalSize;
        this.connectionPoolRemoteSize = connectionPoolRemoteSize;
        this.connectionTimeout = connectionTimeout;
        this.requestTimeout = requestTimeout;
        this.consistencyLevel = consistencyLevel;
        this.sslEnabled = sslEnabled;
        this.sslKeystorePath = sslKeystorePath;
        this.sslKeystorePassword = sslKeystorePassword;
        this.sslTruststorePath = sslTruststorePath;
        this.sslTruststorePassword = sslTruststorePassword;
        this.rowType = rowType;
        this.appendSetFields = appendSetFields != null ? appendSetFields : new java.util.HashSet<>();
        this.primaryKeys = primaryKeys != null ? primaryKeys : new java.util.ArrayList<>();
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        super.open(parameters);
        
        LOG.info("Opening ScyllaDB sink connection to {}", url);
        
        // 创建集群连接
        this.cluster = createCluster();
        this.session = cluster.connect(keySpace);
        
        // 创建序列化器
        this.serializer = new ScyllaDBRowDataSerializer(rowType, primaryKeys);
        
        // 准备 SQL 语句
        prepareStatements();
        
        // 初始化缓冲区
        this.buffer = new ArrayList<>();
        
        // 启动定时刷新任务
        this.scheduler = new ScheduledThreadPoolExecutor(1);
        scheduler.scheduleAtFixedRate(
                this::flushBuffer,
                bufferFlushInterval.toMillis(),
                bufferFlushInterval.toMillis(),
                TimeUnit.MILLISECONDS);
        
        LOG.info("ScyllaDB sink connection opened successfully");
    }

    @Override
    public void invoke(RowData value, Context context) throws Exception {
        synchronized (buffer) {
            buffer.add(value);
            
            if (buffer.size() >= bufferFlushMaxRows) {
                flushBuffer();
            }
        }
    }

    @Override
    public void close() throws Exception {
        LOG.info("Closing ScyllaDB sink connection");
        
        // 刷新剩余的数据
        flushBuffer();
        
        // 关闭调度器
        if (scheduler != null) {
            scheduler.shutdown();
        }
        
        if (session != null) {
            session.close();
        }
        
        if (cluster != null) {
            cluster.close();
        }
        
        super.close();
        LOG.info("ScyllaDB sink connection closed");
    }

    private void flushBuffer() {
        synchronized (buffer) {
            if (buffer.isEmpty()) {
                return;
            }
            
            LOG.debug("Flushing {} rows to ScyllaDB", buffer.size());
            
            try {
                BatchStatement batch = new BatchStatement();
                
                for (RowData row : buffer) {
                    Statement statement = createStatement(row);
                    if (statement != null) {
                        batch.add(statement);
                    }
                }
                
                if (batch.size() > 0) {
                    batch.setConsistencyLevel(ConsistencyLevel.valueOf(consistencyLevel));
                    executeWithRetry(batch);
                }
                
                buffer.clear();
                LOG.debug("Successfully flushed batch to ScyllaDB");
                
            } catch (Exception e) {
                LOG.error("Failed to flush batch to ScyllaDB", e);
                throw new RuntimeException("Failed to flush batch to ScyllaDB", e);
            }
        }
    }

    private Statement createStatement(RowData row) {
        RowKind kind = row.getRowKind();
        
        switch (kind) {
            case INSERT:
            case UPDATE_AFTER:
                // 检查是否需要使用 Set/Array 追加操作
                if (shouldUseAppendOperation(row)) {
                    return createAppendSetStatement(row);
                } else {
                    return serializer.serializeInsert(row, insertStatement);
                }
                
            case DELETE:
                return serializer.serializeDelete(row, deleteStatement);
                
            case UPDATE_BEFORE:
                // 在 ScyllaDB 中通常忽略 UPDATE_BEFORE
                return null;
                
            default:
                LOG.warn("Unsupported row kind: {}", kind);
                return null;
        }
    }

    /**
     * 判断是否应该使用 Set 追加操作
     */
    private boolean shouldUseAppendOperation(RowData row) {
        // 简化实现：如果配置了 appendSetFields，则使用追加操作
        return !appendSetFields.isEmpty();
    }

    /**
     * 创建 Set 追加操作的语句
     */
    private Statement createAppendSetStatement(RowData row) {
        if (appendSetStatement == null) {
            prepareAppendSetStatement();
        }
        return serializer.serializeAppendSet(row, appendSetStatement, appendSetFields);
    }

    private void executeWithRetry(Statement statement) {
        int attempts = 0;
        Exception lastException = null;
        
        while (attempts < maxRetries) {
            try {
                session.execute(statement);
                return;
            } catch (Exception e) {
                lastException = e;
                attempts++;
                
                if (attempts < maxRetries) {
                    LOG.warn("Attempt {} failed, retrying... Error: {}", attempts, e.getMessage());
                    try {
                        Thread.sleep(1000 * attempts);
                    } catch (InterruptedException ie) {
                        Thread.currentThread().interrupt();
                        throw new RuntimeException("Interrupted during retry", ie);
                    }
                }
            }
        }
        
        throw new RuntimeException("Failed to execute after " + maxRetries + " attempts", lastException);
    }

    private Cluster createCluster() {
        String[] hostAndPort = url.split(":");
        String host = hostAndPort[0];
        int port = hostAndPort.length > 1 ? Integer.parseInt(hostAndPort[1]) : 9042;
        
        Cluster.Builder builder = Cluster.builder()
                .addContactPoint(host)
                .withPort(port)
                .withCredentials(username, password)
                .withLoadBalancingPolicy(
                        new TokenAwarePolicy(DCAwareRoundRobinPolicy.builder().build()))
                .withRetryPolicy(DefaultRetryPolicy.INSTANCE)
                .withPoolingOptions(new PoolingOptions()
                        // setConnectionsPerHost 修正参数数量
                        .setConnectionsPerHost(HostDistance.LOCAL, connectionPoolLocalSize, connectionPoolLocalSize)
                        .setConnectionsPerHost(HostDistance.REMOTE, connectionPoolRemoteSize, connectionPoolRemoteSize))
                .withSocketOptions(new SocketOptions()
                        .setConnectTimeoutMillis((int) connectionTimeout.toMillis())
                        .setReadTimeoutMillis((int) requestTimeout.toMillis()));

        // SSL 配置
        if (sslEnabled) {
            try {
                SSLOptions sslOptions = RemoteEndpointAwareJdkSSLOptions.builder()
                        .withSSLContext(createSSLContext())
                        .build();
                builder.withSSL(sslOptions);
            } catch (Exception e) {
                throw new RuntimeException("Failed to configure SSL", e);
            }
        }

        return builder.build();
    }

    private SSLContext createSSLContext() throws Exception {
        // 简化的 SSL 上下文创建，实际使用时需要根据具体需求配置
        return SSLContext.getDefault();
    }

    private void prepareStatements() {
        // 构建 INSERT 语句
        StringBuilder insertSql = new StringBuilder();
        insertSql.append("INSERT INTO ").append(keySpace).append(".").append(tableName).append(" (");
        
        List<String> fieldNames = rowType.getFieldNames();
        for (int i = 0; i < fieldNames.size(); i++) {
            if (i > 0) {
                insertSql.append(", ");
            }
            insertSql.append(fieldNames.get(i));
        }
        
        insertSql.append(") VALUES (");
        for (int i = 0; i < fieldNames.size(); i++) {
            if (i > 0) {
                insertSql.append(", ");
            }
            insertSql.append("?");
        }
        insertSql.append(")");
        
        this.insertStatement = session.prepare(insertSql.toString());
        
        // 构建 DELETE 语句（基于主键）
        StringBuilder deleteSql = new StringBuilder();
        deleteSql.append("DELETE FROM ").append(keySpace).append(".").append(tableName);
        deleteSql.append(" WHERE ");
        
        // 这里需要根据表的主键来构建 WHERE 子句
        // 简化处理，假设第一个字段是主键
        deleteSql.append(fieldNames.get(0)).append(" = ?");
        
        this.deleteStatement = session.prepare(deleteSql.toString());
        
        LOG.debug("Prepared statements: INSERT and DELETE");
    }

    /**
     * 准备 Set/Array 追加操作的语句
     */
    private void prepareAppendSetStatement() {
        if (appendSetFields.isEmpty()) {
            return;
        }
        
        StringBuilder updateSql = new StringBuilder();
        updateSql.append("UPDATE ").append(keySpace).append(".").append(tableName);
        updateSql.append(" SET ");
        
        List<String> fieldNames = rowType.getFieldNames();
        boolean first = true;
        
        for (String fieldName : appendSetFields) {
            if (!first) {
                updateSql.append(", ");
            }
            // 对于 Array/List 类型，使用 = field_name + ? 语法进行追加
            updateSql.append(fieldName).append(" = ").append(fieldName).append(" + ?");
            first = false;
        }
        
        updateSql.append(" WHERE ");
        
        // 构建主键的 WHERE 子句
        // 优先使用显式声明的主键，如果没有则使用除了 appendSetFields 之外的所有字段作为主键
        boolean firstKey = true;
        if (!primaryKeys.isEmpty()) {
            // 使用显式声明的主键
            for (String keyField : primaryKeys) {
                if (!firstKey) {
                    updateSql.append(" AND ");
                }
                updateSql.append(keyField).append(" = ?");
                firstKey = false;
            }
        } else {
            // 回退到使用除了 appendSetFields 之外的所有字段作为主键
            for (String fieldName : fieldNames) {
                if (!appendSetFields.contains(fieldName)) {
                    if (!firstKey) {
                        updateSql.append(" AND ");
                    }
                    updateSql.append(fieldName).append(" = ?");
                    firstKey = false;
                }
            }
        }
        
        this.appendSetStatement = session.prepare(updateSql.toString());
        
        LOG.debug("Prepared append set statement: {}", updateSql.toString());
    }
} 